from ultralytics import YOLO

# Load a model
# model = YOLO('yolov8-SwinTransformer.yaml')  # build a new model from YAML
# model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)
# model = YOLO('yolov8-SwinTransformer.yaml').load('yolov8n.pt')  # build from YAML and transfer weights

# model = YOLO('yolov8.yaml')  # build a new model from YAML
# model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)
model = YOLO('yolov8.yaml').load('yolov8n.pt')  # build from YAML and transfer weights

# model = YOLO('yolov8n.yaml')  # build a new model from YAML
# model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)
# model = YOLO('yolov8n.yaml').load('yolov8n.pt')  # build from YAML and transfer weights
if __name__ == '__main__':
    # Train the model
    model.train(data='./pet.yaml', epochs=200, imgsz=640, batch=32)